(* image_functions.ml - Somme functions applyed on images
 * Licence: LGPLv3
 * Authors: Franck "kushou" Michea < michea_f@epita.fr >
 *)

let fst3 (x,_,_) = x
let ( |> ) x f = f x
  
let iter f img =
  for x = 0 to img#get_width - 1 do
    for y = 0 to img#get_height - 1 do f (x,y) img done
  done

let main img img_t init f =
  if img#exists img_t then ()
  else begin
    img#create img_t;
    init ();
    iter f img;
    img#save img_t
  end
    
let grey_max = 230
let grey_min = 25

let level (r,g,b)=
  let op i = float_of_int i /. 255.
  in 0.3 *. op r +. 0.59 *. op g +. 0.11 *. op b

let color2grey rgb =
  let op = int_of_float ((level rgb) *. 255.)
  in (op,op,op)

let get_color_number color =
  let rec sub n = function
    | [] -> -1
    | x::_ when fst x = color -> n
    | _::xs -> sub (n+1) xs
  in sub 0

let image2grey img =
  let init () = () in
  let f (x,y) img =
    let img_ = img#get Types.Grey in
    let color = color2grey (Sdlvideo.get_pixel_color img_ x y)
    in Sdlvideo.put_pixel_color img_ x y color
  in main img Types.Grey init f

(* Canny *)

let noise_filter =
  [|[| 2; 4; 5; 4; 2|];
    [| 4; 9;12; 9; 4|];
    [| 5;12;15;12; 5|];
    [| 4; 9;12; 9; 4|];
    [| 2; 4; 5; 4; 2|]|]

let g_filter = [|-1;0;1|]

let coef = 155.

let get_color img img_t x y =
  let (x_, y_) = (ref x, ref y) in
  begin
    if x < 0 then x_ := -x
    else if img#get_height <= x then x_ := 2 * img#get_height - x;
    if y < 0 then y_ := -y
    else if img#get_width <= y then y_ := 2 * img#get_width - y;
    Sdlvideo.get_pixel_color (img#get img_t) (!x_) (!y_) |> fst3
  end

let noise_reduction img =
  let init () = () in
  let f (x,y) img =
    let sum = ref 0 in
    for i = 0 to 4 do
      for j = 0 to 4 do
        let color = get_color img Types.Base (x + i - 2) (y + j - 2)
        in sum := !sum + color * noise_filter.(i).(j)
      done
    done;
    let c = int_of_float (1. /. coef *. float_of_int (!sum))
    in Sdlvideo.put_pixel_color (img#get Types.NReduc) x y (c,c,c)
  in main img Types.NReduc init f

let edge_detection img =
  let init () = image2grey img in
  let f (x,y) img =
    let (gx, gy, g) = (ref 0, ref 0, ref 0) in
    for i = 0 to 2 do
      gx := !gx + get_color img Types.Grey (x + i - 1) y * g_filter.(i);
      gy := !gy + get_color img Types.Grey x (y + i - 1) * g_filter.(i);
    done;
    g := abs (!gx) + abs (!gy);
    if !g <> 0 then Sdlvideo.put_pixel_color (img#get Types.Edges) x y (0,0,0)
    else Sdlvideo.put_pixel_color (img#get Types.Edges) x y (255,255,255);
  in main img Types.Edges init f

(* End Of Canny *)

let blit_edges img img_t =
  let init () = edge_detection img in
  let f (x,y) img =
    let c_e = Sdlvideo.get_pixel_color (img#get Types.Edges) x y in
    if c_e = Sdlvideo.black then
      Sdlvideo.put_pixel_color (img#get img_t) x y c_e
  in main img img_t init f

let iwedges img = blit_edges img Types.IWEdges
let gwedges img = blit_edges img Types.GWEdges

let rec add_coords rgb coords = function
  | [] -> (rgb,coords::[])::[]
  | (color,a)::t ->
    if color = rgb then (color,coords::a)::t
    else (color,a)::(add_coords rgb coords t)

let detect_colors img =
  if img#get_colors <> [] then ()
  else
    let init () = () in
    let f (x,y) img =
      let color = Sdlvideo.get_pixel_color (img#get Types.Base) x y in
      if color <> Sdlvideo.black then
	img#set_colors (add_coords color (x,y) img#get_colors)
    in main img Types.RH init f

let highlight_color img img_tmp (rgb, zone) =
  let f (x,y) =
    if Sdlvideo.get_pixel_color img_tmp x y <> Sdlvideo.black then
      Sdlvideo.put_pixel_color img_tmp x y rgb
  in List.iter f zone

(* Heights Manager *)

let cur_nb = ref 0
let cur_file = ref ""
let cur_color = ref ((0,0,0),([]:Types.coordlist))

let get_hfile () = !cur_file
let get_nb () = !cur_nb

let init_heights img =
  let () =
    gwedges img;
    detect_colors img in
  let rec sub n = function
    | [] -> ()
    | x::xs ->
      let img_ = img#get_copy Types.GWEdges
      in begin
	highlight_color img img_ x;
	Sdlvideo.save_BMP img_ (img#get_img_zone_name n);
	sub (n+1) xs end
  in sub 0 img#get_colors

let get_next img =
  if !cur_nb < img#get_colors_length then begin
    cur_color := img#get_color !cur_nb;
    cur_file := img#get_img_zone_name !cur_nb;
    cur_nb := !cur_nb + 1;
    true end
  else false

let end_heights img =
  for i = 0 to img#get_colors_length - 1 do
    Unix.unlink (img#get_img_zone_name i)
  done

let remove_tmp img =
  Unix.unlink (img#get_filename^".edges.bmp");
  Unix.unlink (img#get_filename^".grey.bmp");
  Unix.unlink (img#get_filename^".gwedges.bmp");
  Unix.unlink (img#get_filename^".iwedges.bmp");
  Unix.unlink (img#get_filename^".bmp");
  Unix.unlink (img#get_filename^".heights.txt")

let register_height img h =
  let (a,b) = !cur_color in img#push_heights (h, a, b)

let save_heights img =
  if (List.length img#get_colors) <> (List.length img#get_heights) then
    Printf.printf "All zone's heights haven't been set.\n"
  else Io.write_heights img

let request_heights img =
  init_heights img;
  while get_next img do
    Io.show img (Sdlloader.load_image (get_hfile ())) false;
    register_height img (Io.request_height (get_nb () - 1))
  done;
  save_heights img;
  end_heights img

let set_pas w h =
  let rec get_int () =
    Printf.printf "Please enter the pas for the grid : ";
    try read_int () with
	Failure _ -> get_int ()
      | End_of_file -> exit 0 in
  let coef = get_int () in
  let (mn,mx) = (min w h, max w h) in
  let pas = ref (-1) in
  let () = for i = coef to mn do
      if !pas <> -1 then
	if mn mod i = 0 && mx mod i = 0 then pas := i
	else pas := -1
      else pas := -1
    done
  in if !pas = -1 then coef else !pas

